import json
import os
import requests
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class LLMClient:
    def __init__(self, model_config=None):
        """
        初始化大模型客户端
        
        Args:
            model_config (dict): 大模型配置参数
        """
        # 从环境变量读取配置
        env_config = {
            "api_key": os.environ.get("LLM_TOKEN", "your_llm_token"),
            "api_base": os.environ.get("LLM_BASE_URL", "https://api-inference.modelscope.cn/v1/"),
            "model": os.environ.get("LLM_MODEL_NAME", "Qwen/Qwen3-Coder-480B-A35B-Instruct")
        }
        
        # 默认配置
        self.default_config = {
            "api_key": env_config["api_key"],
            "api_base": env_config["api_base"],
            "model": env_config["model"],
            "temperature": 0.7,
            "max_tokens": 2000
        }
        
        # 合并用户配置
        self.config = self.default_config.copy()
        if model_config:
            self.config.update(model_config)
        
        # 初始化会话状态
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {self.config['api_key']}",
            "Content-Type": "application/json"
        })
        
        logger.info("LLM客户端初始化完成")
    
    def _prepare_messages(self, user_input, system_prompt, tools=None):
        """
        准备发送给大模型的消息
        
        Args:
            user_input (str): 用户输入
            system_prompt (str): 系统提示
            tools (list): 可用工具列表
            
        Returns:
            list: 消息列表
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_input}
        ]
        return messages
    
    def _prepare_tool_params(self, tools):
        """
        准备工具参数，转换为MCP协议格式
        
        Args:
            tools (list): 工具信息列表
            
        Returns:
            list: 格式化后的工具列表
        """
        formatted_tools = []
        for tool in tools:
            # 确保参数格式正确，避免JSON序列化错误
            try:
                # 验证parameters是否为有效的JSON结构
                parameters = tool.get("parameters", {})
                # 尝试JSON序列化以验证格式
                json.dumps(parameters)
            except (TypeError, ValueError):
                logger.warning(f"工具 {tool.get('name')} 的参数格式无效，使用空字典")
                parameters = {}
            
            formatted_tool = {
                "type": "function",
                "function": {
                    "name": tool.get("name", ""),
                    "description": tool.get("description", ""),
                    "parameters": parameters
                }
            }
            formatted_tools.append(formatted_tool)
        return formatted_tools
    
    def call_with_tools(self, tools, user_input, system_prompt):
        """
        使用工具回调方式访问大模型
        
        Args:
            tools (list): 需要回调的工具信息
            user_input (str): 用户输入信息
            system_prompt (str): 系统提示
            
        Returns:
            dict: 包含调用结果和详细信息
        """
        try:
            # 验证输入参数
            if not isinstance(user_input, str):
                user_input = str(user_input) if user_input is not None else ""
            if not isinstance(system_prompt, str):
                system_prompt = str(system_prompt) if system_prompt is not None else ""
            
            # 准备请求数据
            data = {
                "model": self.config["model"],
                "messages": self._prepare_messages(user_input, system_prompt),
                "tools": self._prepare_tool_params(tools),
                "tool_choice": "auto",  # 让模型自动决定是否调用工具
                "temperature": self.config["temperature"],
                "max_tokens": self.config["max_tokens"]
            }
            
            # 再次验证数据的JSON格式
            json.dumps(data)
            
            logger.info(f"准备调用大模型，工具数量: {len(tools)}")
            
            # 打印大模型输入数据（截断以避免日志过大）
            request_data_str = json.dumps(data, ensure_ascii=False)
            logger.info(f"大模型输入数据: {request_data_str[:500]}...")
            
            # 发送请求 - 确保URL正确拼接，避免多余的斜杠
            api_base = self.config['api_base'].rstrip('/')
            response = self.session.post(
                f"{api_base}/chat/completions",
                json=data
            )
            
            response.raise_for_status()
            result = response.json()
            
            # 打印大模型输出数据（截断以避免日志过大）
            response_data_str = json.dumps(result, ensure_ascii=False)
            logger.info(f"大模型输出数据: {response_data_str[:500]}...")
            
            # 处理响应
            if "choices" in result and len(result["choices"]) > 0:
                choice = result["choices"][0]
                
                # 检查是否调用了工具
                if "tool_calls" in choice.get("message", {}) and choice["message"]["tool_calls"]:
                    tool_call = choice["message"]["tool_calls"][0]
                    if tool_call.get("type") == "function":
                        # 打印工具调用详情
                        tool_call_str = json.dumps(tool_call, ensure_ascii=False)
                        logger.info(f"工具调用详情: {tool_call_str}")
                        
                        # 获取原始arguments字符串
                        raw_arguments = tool_call["function"].get("arguments", "{}")
                        logger.info(f"原始工具参数字符串: {raw_arguments}")
                        
                        # 安全地解析工具参数
                        try:
                            function_args = json.loads(raw_arguments)
                            logger.info(f"成功解析工具参数: {function_args}")
                        except json.JSONDecodeError as e:
                            logger.error(f"解析工具参数失败: {str(e)}")
                            logger.error(f"问题参数内容: '{raw_arguments[:200]}...'")
                            
                            # 增强的修复机制
                            try:
                                # 方法1: 清理并修复JSON字符串
                                # 1. 移除开头和结尾的可能的额外引号
                                fixed_args = raw_arguments
                                if fixed_args.startswith('"') and fixed_args.endswith('"'):
                                    fixed_args = fixed_args[1:-1]
                                
                                # 2. 保留原始换行符，不进行转义（避免将换行符转为文本）
                                # fixed_args = fixed_args.replace('\n', '\\n') - 这行导致换行符被转义为文本
                                
                                # 3. 修复嵌套引号问题
                                fixed_args = fixed_args.replace('""', '"')
                                fixed_args = fixed_args.replace('\\"', '\\\\"')
                                
                                # 4. 确保JSON结构完整
                                if not fixed_args.endswith('}'):
                                    # 尝试找到最后一个有效的闭合括号位置
                                    last_brace_pos = fixed_args.rfind('}')
                                    if last_brace_pos != -1:
                                        fixed_args = fixed_args[:last_brace_pos+1]
                                    else:
                                        # 如果没有找到闭合括号，添加一个
                                        fixed_args += '}'
                                
                                function_args = json.loads(fixed_args)
                                logger.info(f"修复后成功解析工具参数")
                            except:
                                # 方法2: 使用更宽松的解析
                                try:
                                    # 尝试简单的键值对解析
                                    from ast import literal_eval
                                    # 替换可能的问题字符
                                    safe_args = raw_arguments.replace('\\', '').replace('\n', '')
                                    function_args = literal_eval(safe_args)
                                    logger.info(f"使用literal_eval成功解析工具参数")
                                except:
                                    # 方法3: 提取content字段的简单方法
                                    try:
                                        # 直接提取content的值
                                        content_start = raw_arguments.find('"content": "') + len('"content": "')
                                        content_end = raw_arguments.rfind('"')
                                        if content_start > 0 and content_end > content_start:
                                            content = raw_arguments[content_start:content_end]
                                            # 安全地将\n替换为实际换行符
                                            content = content.replace('\\n', '\n')
                                            function_args = {"content": content}
                                            logger.info(f"成功提取并处理content字段中的换行符")
                                        else:
                                            # 使用空参数继续执行
                                            function_args = {}
                                            logger.warning("使用空参数继续执行")
                                    except:
                                        # 最终 fallback
                                        function_args = {}
                                        logger.warning("所有修复方法失败，使用空参数继续执行")
                        
                    logger.info(f"大模型选择调用工具: {tool_call['function'].get('name', 'unknown')}")
                    
                    return {
                        "status": "success",
                        "tool_call": {
                            "name": tool_call["function"].get("name", "unknown"),
                            "parameters": function_args,
                            "raw_arguments": raw_arguments
                        },
                        "raw_response": result
                    }
                else:
                    # 直接回答
                    content = choice.get("message", {}).get("content", "")
                    logger.info("大模型直接回答用户")
                    return {
                        "status": "success",
                        "content": content,
                        "raw_response": result
                    }
            
            return {"status": "success", "raw_response": result}
            
        except json.JSONDecodeError as e:
            logger.error(f"JSON格式错误: {str(e)}")
            # 如果是响应解析错误，打印响应内容
            if 'response' in locals():
                logger.error(f"原始响应内容: {response.text[:500]}...")
            return {
                "status": "error",
                "error_type": "json_error",
                "error_message": f"JSON格式错误: {str(e)}"
            }
        except requests.exceptions.RequestException as e:
            logger.error(f"API请求失败: {str(e)}")
            # 如果是请求失败，打印响应内容（如果有）
            if 'response' in locals() and response.content:
                logger.error(f"错误响应内容: {response.text[:500]}...")
            return {
                "status": "error",
                "error_type": "api_error",
                "error_message": str(e)
            }
        except Exception as e:
            logger.error(f"调用大模型失败: {str(e)}")
            import traceback
            logger.error(f"异常堆栈: {traceback.format_exc()}")
            return {
                "status": "error",
                "error_type": "unknown_error",
                "error_message": str(e)
            }
    
    def set_api_key(self, api_key):
        """
        设置API密钥
        
        Args:
            api_key (str): 新的API密钥
        """
        self.config["api_key"] = api_key
        self.session.headers.update({
            "Authorization": f"Bearer {api_key}"
        })
        logger.info("API密钥已更新")
    
    def close(self):
        """
        关闭会话
        """
        self.session.close()
        logger.info("LLM客户端已关闭")

# 便捷函数
def create_llm_client(model_config=None):
    """
    创建LLM客户端实例
    
    Args:
        model_config (dict): 模型配置
        
    Returns:
        LLMClient: 客户端实例
    """
    return LLMClient(model_config)

def call_llm_with_tools(tools, user_input, system_prompt, model_config=None):
    """
    便捷函数：使用工具回调方式调用大模型
    
    Args:
        tools (list): 工具信息
        user_input (str): 用户输入
        system_prompt (str): 系统提示
        model_config (dict): 模型配置
        
    Returns:
        dict: 调用结果
    """
    client = create_llm_client(model_config)
    try:
        return client.call_with_tools(tools, user_input, system_prompt)
    finally:
        client.close()